import cv2
import mediapipe as mp
import numpy as np
import time
from PIL import Image, ImageDraw, ImageFont
from sklearn.tree import DecisionTreeClassifier
import pickle
import os

# 初始化手部检测模型
mp_hands = mp.solutions.hands
hands = mp_hands.Hands(
    max_num_hands=1,
    min_detection_confidence=0.6,
    min_tracking_confidence=0.6,
    model_complexity=1
)

# 手势映射表
GESTURE_MAP = {0: "石头", 1: "剪刀", 2: "布", -1: "未识别", -2: "手背"}
# 必胜策略
WINNING_STRATEGY = {0: 2, 1: 0, 2: 1}

# 游戏状态常量
GAME_STATE_READY = 0
GAME_STATE_COUNTDOWN = 1
GAME_STATE_PLAYING = 2
GAME_STATE_RESULT = 3
GAME_STATE_RESTART = 4


# 解决OpenCV中文显示问题的函数
def cv2_put_chinese_text(img, text, position, font_size=24, color=(0, 255, 0)):
    try:
        rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        pil_img = Image.fromarray(rgb_img)
        draw = ImageDraw.Draw(pil_img)
        try:
            font = ImageFont.truetype("simhei.ttf", font_size)
        except IOError:
            try:
                font = ImageFont.load_default()
                print("警告：未找到中文字体，使用默认字体")
            except Exception as e:
                print(f"错误：无法加载任何字体 - {e}")
                return img

        draw.text(position, text, font=font, fill=color)
        return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    except Exception as e:
        print(f"绘制文字时出错: {e}")
        return img


# 加载手势图片
def load_gesture_images():
    try:
        rock_img = cv2.imread("rock.png")
        scissors_img = cv2.imread("scissors.png")
        paper_img = cv2.imread("paper.png")

        if rock_img is not None:
            rock_img = cv2.resize(rock_img, (200, 200))
        if scissors_img is not None:
            scissors_img = cv2.resize(scissors_img, (200, 200))
        if paper_img is not None:
            paper_img = cv2.resize(paper_img, (200, 200))

        return {0: rock_img, 1: scissors_img, 2: paper_img}
    except Exception as e:
        print(f"加载手势图片时出错: {e}")
        return None


# 从手部关键点提取特征
def extract_features(hand_landmarks):
    landmarks = hand_landmarks.landmark
    features = []

    # 计算关键点之间的相对距离和角度作为特征
    wrist = (landmarks[mp_hands.HandLandmark.WRIST].x, landmarks[mp_hands.HandLandmark.WRIST].y)

    for i in range(1, 21):  # 跳过手腕(0)
        point = (landmarks[i].x, landmarks[i].y)
        # 计算相对于手腕的位置
        features.append(point[0] - wrist[0])
        features.append(point[1] - wrist[1])

    # 计算手指之间的相对位置
    finger_tips = [mp_hands.HandLandmark.THUMB_TIP,
                   mp_hands.HandLandmark.INDEX_FINGER_TIP,
                   mp_hands.HandLandmark.MIDDLE_FINGER_TIP,
                   mp_hands.HandLandmark.RING_FINGER_TIP,
                   mp_hands.HandLandmark.PINKY_TIP]

    for i in range(len(finger_tips)):
        for j in range(i + 1, len(finger_tips)):
            x1, y1 = landmarks[finger_tips[i]].x, landmarks[finger_tips[i]].y
            x2, y2 = landmarks[finger_tips[j]].x, landmarks[finger_tips[j]].y
            # 计算距离
            dist = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            features.append(dist)

    return features


# 数据收集函数
def collect_data():
    print("开始收集训练数据...")
    print("请按以下顺序做出手势：石头(0) -> 剪刀(1) -> 布(2)")
    print("每种手势收集100帧数据，按任意键开始...")
    cv2.waitKey(0)

    cap = cv2.VideoCapture(0)
    data = []
    labels = []

    for gesture_id in [0, 1, 2]:  # 石头、剪刀、布
        print(f"\n请做出 {GESTURE_MAP[gesture_id]} 的手势")
        print("3秒后开始收集数据...")
        time.sleep(3)

        count = 0
        while count < 100:  # 每种手势收集100帧
            success, frame = cap.read()
            if not success:
                print("无法获取图像")
                break

            frame = cv2.flip(frame, 1)
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = hands.process(rgb_frame)

            if results.multi_hand_landmarks:
                hand_landmarks = results.multi_hand_landmarks[0]
                # 绘制手部关键点
                mp.solutions.drawing_utils.draw_landmarks(
                    frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)

                # 提取特征
                features = extract_features(hand_landmarks)
                data.append(features)
                labels.append(gesture_id)

                count += 1
                frame = cv2_put_chinese_text(frame, f"收集进度: {count}/100", (10, 30), 24)

            frame = cv2_put_chinese_text(frame, f"请做出 {GESTURE_MAP[gesture_id]} 的手势", (10, 70), 24)
            cv2.imshow("收集训练数据", frame)

            if cv2.waitKey(1) == 27:  # 按ESC退出
                cap.release()
                cv2.destroyAllWindows()
                return None, None

        print(f"{GESTURE_MAP[gesture_id]} 手势数据收集完成")

    cap.release()
    cv2.destroyAllWindows()

    return np.array(data), np.array(labels)


# 训练决策树模型
def train_model(data, labels):
    print("开始训练模型...")
    model = DecisionTreeClassifier(max_depth=5, random_state=42)
    model.fit(data, labels)
    print("模型训练完成")

    # 保存模型
    with open('gesture_model.pkl', 'wb') as f:
        pickle.dump(model, f)
    print("模型已保存到 gesture_model.pkl")

    return model


# 加载预训练模型
def load_model():
    if os.path.exists('gesture_model.pkl'):
        try:
            with open('gesture_model.pkl', 'rb') as f:
                model = pickle.load(f)
            print("已加载预训练模型")
            return model
        except:
            print("无法加载预训练模型，将重新训练")
    return None


# 判断胜负
def judge_winner(player_gesture, bot_gesture):
    if player_gesture == bot_gesture:
        return 0  # 平局
    elif (player_gesture == 0 and bot_gesture == 2) or \
            (player_gesture == 1 and bot_gesture == 0) or \
            (player_gesture == 2 and bot_gesture == 1):
        return -1  # 玩家输
    else:
        return 1  # 玩家赢


# 在画面上显示结果
def display_result(frame, player_gesture, bot_gesture, gesture_images, game_state, winner=None, countdown=None):
    height, width, _ = frame.shape

    if game_state == GAME_STATE_READY:
        frame = cv2_put_chinese_text(frame, "准备开始游戏...", (width // 2 - 150, height // 2 - 15), 30, (0, 255, 255))
        frame = cv2_put_chinese_text(frame, "请将手放入画面中", (width // 2 - 150, height // 2 + 30), 24, (0, 255, 255))
        frame = cv2_put_chinese_text(frame, "按任意键开始", (width // 2 - 100, height // 2 + 75), 24, (0, 255, 255))

    elif game_state == GAME_STATE_COUNTDOWN:
        frame = cv2_put_chinese_text(frame, f"准备: {countdown}", (width // 2 - 50, height // 2 - 15), 60, (0, 0, 255))

    elif game_state == GAME_STATE_PLAYING:
        if player_gesture == -2:
            frame = cv2_put_chinese_text(frame, "检测到手背，请展示手掌", (10, 30), 24, (0, 0, 255))
        else:
            frame = cv2_put_chinese_text(frame, f"你: {GESTURE_MAP[player_gesture]}", (10, 30), 24, (0, 255, 0))

        if player_gesture in [0, 1, 2]:
            if gesture_images and gesture_images.get(bot_gesture) is not None:
                bot_img = gesture_images[bot_gesture]
                img_height, img_width, _ = bot_img.shape
                frame[10:10 + img_height, width - img_width - 10:width - 10] = bot_img
            else:
                frame = cv2_put_chinese_text(frame, f"机器: {GESTURE_MAP[bot_gesture]}",
                                             (width - 200, 30), 24, (0, 0, 255))

    elif game_state == GAME_STATE_RESULT:
        frame = cv2_put_chinese_text(frame, f"你: {GESTURE_MAP[player_gesture]}", (10, 30), 24, (0, 255, 0))

        if gesture_images and gesture_images.get(bot_gesture) is not None:
            bot_img = gesture_images[bot_gesture]
            img_height, img_width, _ = bot_img.shape
            frame[10:10 + img_height, width - img_width - 10:width - 10] = bot_img
        else:
            frame = cv2_put_chinese_text(frame, f"机器: {GESTURE_MAP[bot_gesture]}",
                                         (width - 200, 30), 24, (0, 0, 255))

        if winner == 0:
            result_text = "平局!"
            color = (0, 255, 255)
        elif winner == 1:
            result_text = "你赢了!"
            color = (0, 255, 0)
        else:
            result_text = "机器赢了!"
            color = (0, 0, 255)

        frame = cv2_put_chinese_text(frame, result_text, (width // 2 - 100, height // 2 - 15), 36, color)

    elif game_state == GAME_STATE_RESTART:
        if winner == 0:
            result_text = "平局!"
            color = (0, 255, 255)
        elif winner == 1:
            result_text = "你赢了!"
            color = (0, 255, 0)
        else:
            result_text = "机器赢了!"
            color = (0, 0, 255)

        frame = cv2_put_chinese_text(frame, result_text, (width // 2 - 100, height // 2 - 60), 36, color)
        frame = cv2_put_chinese_text(frame, "按 R 键重新开始", (width // 2 - 120, height // 2 + 10), 30, (255, 255, 0))
        frame = cv2_put_chinese_text(frame, "按 ESC 键退出", (width // 2 - 120, height // 2 + 60), 30, (255, 255, 0))
        frame = cv2_put_chinese_text(frame, "按 T 键训练新模型", (width // 2 - 120, height // 2 + 110), 30,
                                     (255, 165, 0))

    return frame


def main():
    # 尝试加载预训练模型
    model = load_model()

    # 如果没有预训练模型，收集数据并训练
    if model is None:
        data, labels = collect_data()
        if data is not None and len(data) > 0:
            model = train_model(data, labels)
        else:
            print("没有收集到训练数据，将使用基于规则的识别")

    # 加载手势图片
    gesture_images = load_gesture_images()

    # 打开摄像头
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    # 初始化游戏状态
    game_state = GAME_STATE_READY
    current_player_gesture = -1
    current_bot_gesture = -1
    last_gesture_time = 0
    gesture_stability_threshold = 0.5
    last_detected_gesture = -1
    countdown_start_time = 0
    countdown_duration = 5
    result_display_time = 3
    game_winner = 0

    cv2.namedWindow("猜拳游戏 - 按ESC退出", cv2.WINDOW_NORMAL)

    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break

        frame = cv2.flip(frame, 1)
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        current_time = time.time()

        # 处理按键
        key = cv2.waitKey(1)
        if key == 27:  # ESC
            break
        elif key == ord('t') or key == ord('T'):  # T键重新训练模型
            data, labels = collect_data()
            if data is not None and len(data) > 0:
                model = train_model(data, labels)
            else:
                print("没有收集到足够的训练数据")

        if game_state == GAME_STATE_READY:
            if key != -1 and key != 27 and key != ord('t') and key != ord('T'):
                game_state = GAME_STATE_COUNTDOWN
                countdown_start_time = current_time

        elif game_state == GAME_STATE_COUNTDOWN:
            elapsed_time = current_time - countdown_start_time
            remaining_time = countdown_duration - int(elapsed_time)

            if remaining_time > 0:
                frame = display_result(frame, current_player_gesture, current_bot_gesture,
                                       gesture_images, game_state, countdown=remaining_time)
            else:
                game_state = GAME_STATE_PLAYING
                last_gesture_time = current_time

        elif game_state == GAME_STATE_PLAYING:
            results = hands.process(rgb_frame)
            detected_gesture = -1

            if results.multi_hand_landmarks and results.multi_handedness:
                for hand_landmarks, hand_info in zip(results.multi_hand_landmarks, results.multi_handedness):
                    mp.solutions.drawing_utils.draw_landmarks(
                        frame, hand_landmarks, mp_hands.HAND_CONNECTIONS)

                    # 提取特征
                    features = extract_features(hand_landmarks)

                    # 使用模型预测手势
                    if model is not None and len(features) > 0:
                        try:
                            features_array = np.array(features).reshape(1, -1)
                            detected_gesture = model.predict(features_array)[0]
                        except Exception as e:
                            print(f"模型预测出错: {e}")
                            detected_gesture = -1
                    else:
                        # 回退到基于规则的识别
                        detected_gesture = recognize_gesture(hand_landmarks, hand_info)

                    if detected_gesture in [0, 1, 2]:
                        if detected_gesture == last_detected_gesture and \
                                current_time - last_gesture_time > gesture_stability_threshold:
                            if detected_gesture != current_player_gesture:
                                current_player_gesture = detected_gesture
                                current_bot_gesture = WINNING_STRATEGY[detected_gesture]
                                game_winner = judge_winner(current_player_gesture, current_bot_gesture)
                                print(
                                    f"用户出: {GESTURE_MAP[detected_gesture]}, 机器出: {GESTURE_MAP[current_bot_gesture]}, 结果: {game_winner}")

                                game_state = GAME_STATE_RESULT
                                result_start_time = current_time
                        else:
                            if detected_gesture != last_detected_gesture:
                                last_gesture_time = current_time

                        last_detected_gesture = detected_gesture

            frame = display_result(frame, current_player_gesture, current_bot_gesture,
                                   gesture_images, game_state)

        elif game_state == GAME_STATE_RESULT:
            if current_time - result_start_time > result_display_time:
                game_state = GAME_STATE_RESTART
            else:
                frame = display_result(frame, current_player_gesture, current_bot_gesture,
                                       gesture_images, game_state, game_winner)

        elif game_state == GAME_STATE_RESTART:
            frame = display_result(frame, current_player_gesture, current_bot_gesture,
                                   gesture_images, game_state, game_winner)

            if key == ord('r') or key == ord('R'):
                game_state = GAME_STATE_READY
                current_player_gesture = -1
                current_bot_gesture = -1
                last_detected_gesture = -1

        cv2.imshow("猜拳游戏 - 按ESC退出", frame)

    cap.release()
    cv2.destroyAllWindows()


# 原有的基于规则的手势识别函数（作为后备）
def recognize_gesture(hand_landmarks, hand_info):
    landmarks = hand_landmarks.landmark

    thumb_open = landmarks[mp_hands.HandLandmark.THUMB_TIP].x > landmarks[mp_hands.HandLandmark.THUMB_IP].x
    index_open = landmarks[mp_hands.HandLandmark.INDEX_FINGER_TIP].y < landmarks[
        mp_hands.HandLandmark.INDEX_FINGER_PIP].y
    middle_open = landmarks[mp_hands.HandLandmark.MIDDLE_FINGER_TIP].y < landmarks[
        mp_hands.HandLandmark.MIDDLE_FINGER_PIP].y
    ring_open = landmarks[mp_hands.HandLandmark.RING_FINGER_TIP].y < landmarks[mp_hands.HandLandmark.RING_FINGER_PIP].y
    pinky_open = landmarks[mp_hands.HandLandmark.PINKY_TIP].y < landmarks[mp_hands.HandLandmark.PINKY_PIP].y

    fingers_open = sum([thumb_open, index_open, middle_open, ring_open, pinky_open])

    is_back_of_hand = landmarks[mp_hands.HandLandmark.WRIST].z < -0.05

    if fingers_open == 0:
        return 0
    elif fingers_open == 2 and index_open and middle_open:
        return 1
    elif fingers_open == 5:
        return 2
    elif is_back_of_hand:
        return -2
    else:
        return -1


if __name__ == "__main__":
    main()